# -*- coding: utf-8 -*-
import os
import sys
from typing import Collection

import numpy as np
import pandas as pd
import math
import pymongo
import torch
import re
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from sklearn.metrics import accuracy_score
from tqdm import tqdm

matplotlib.use("Agg")


if __name__ == "__main__":
    client = pymongo.MongoClient("mongodb://127.0.0.1:27017/")
    db = client["Power_Fault"]
    col_sour = db["data_sour"]
    # all ID of the data, split data by ID
    rect_filter = {'location': {
        '$regex': re.compile(r"(?si:.*整流.*)")}}
    rect_ids = col_sour.distinct("_id",filter=rect_filter)
    inver_ids = col_sour.distinct("_id",filter={'location': {
        '$regex': re.compile(r"(?si:.*逆变.*)")}})
    ids = np.append(rect_ids, inver_ids)
    experimental_len = int(len(ids) * 0.8)
    ids_all = col_sour.distinct("_id")
    print(len(ids_all))
    experimental_len_left = len(ids) - experimental_len
    training_count = int((len(ids_all) - len(ids) * 0.2) * 0.6)
    validation_count = int((len(ids_all) - len(ids) * 0.2) * 0.2)
    test_count = len(ids_all) - training_count - validation_count
    simulation_size = len(ids_all) - len(ids)
    actual_exper_size = len(ids)

    print(f"the size of experimental test set is {experimental_len}") 
    print(f"the size of train set is {training_count}") 
    print(f"the size of validation  set is {validation_count}") 
    print(f"the size of test  set is {test_count}") 
    print(f"the size of actual experiment  set is {actual_exper_size}") 
    print(f"the size of actual simulation  set is {simulation_size}") 
